#!/bin/bash

# 大规模构型生成脚本 - Be原子薄膜（删除顶层原子）
# 生成1500个不同原子排列的构型

# 设置ABACUS路径和运行参数
ABACUS_CMD="/usr/local/bin/abacus"
CORE=10

# 创建结果目录
mkdir -p config_results_1500
RESULT_FILE="config_energy_1500.dat"
echo "# Config_ID Energy Volume Stress_XX Stress_YY Stress_ZZ" > $RESULT_FILE

# 原始STRU文件内容（删除第25个原子）
ORIGINAL_STRU=$(cat << 'EOF'
ATOMIC_SPECIES
Be    9.012  Be_ONCV_PBE-1.0.upf

NUMERICAL_ORBITAL
Be_gga_10au_100Ry_6s2p.orb

LATTICE_CONSTANT
1.889726

LATTICE_VECTORS
11.11980000000  0.000000000000  0.000000000000
0.000000000000  9.630000000000  0.000000000000
0.000000000000  0.000000000000  20.00000000000

ATOMIC_POSITIONS
Direct

Be
0.000
24
0.250000000000  0.166667000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.666667000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.000000000000  0.000000000000  1  1  1  mag  0.0
0.250000000000  0.500000000000  0.000000000000  1  1  1  mag  0.0
0.000000000000  0.333333000000  0.000000000000  1  1  1  mag  0.0
0.250000000000  0.833333000000  0.000000000000  1  1  1  mag  0.0
0.333333000000  0.000000000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.500000000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.166667000000  0.130325000000  1  1  1  mag  0.0
0.333333000000  0.666667000000  0.130325000000  1  1  1  mag  0.0
0.333333000000  0.333333000000  0.130325000000  1  1  1  mag  0.0
0.083333000000  0.833333000000  0.130325000000  1  1  1  mag  0.0
0.750000000000  0.166667000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.666667000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.000000000000  0.000000000000  1  1  1  mag  0.0
0.750000000000  0.500000000000  0.000000000000  1  1  1  mag  0.0
0.500000000000  0.333333000000  0.000000000000  1  1  1  mag  0.0
0.750000000000  0.833333000000  0.000000000000  1  1  1  mag  0.0
0.833333000000  0.000000000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.500000000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.166667000000  0.130325000000  1  1  1  mag  0.0
0.833333000000  0.666667000000  0.130325000000  1  1  1  mag  0.0
0.833333000000  0.333333000000  0.130325000000  1  1  1  mag  0.0
0.583333000000  0.833333000000  0.130325000000  1  1  1  mag  0.0
EOF
)

# 函数：计算两点之间的距离
calculate_distance() {
    local x1=$1 y1=$2 z1=$3
    local x2=$4 y2=$5 z2=$6
    local dx=$(echo "scale=10; $x1 - $x2" | bc -l)
    local dy=$(echo "scale=10; $y1 - $y2" | bc -l)
    local dz=$(echo "scale=10; $z1 - $z2" | bc -l)
    echo "scale=10; sqrt($dx*$dx + $dy*$dy + $dz*$dz)" | bc -l
}

# 函数：检查原子距离是否满足要求
check_atomic_distances() {
    local atoms=("$@")
    local num_atoms=${#atoms[@]}
    local min_distance=1000
    
    for ((i=0; i<num_atoms; i++)); do
        for ((j=i+1; j<num_atoms; j++)); do
            local atom_i=(${atoms[i]})
            local atom_j=(${atoms[j]})
            
            local dist=$(calculate_distance ${atom_i[0]} ${atom_i[1]} ${atom_i[2]} ${atom_j[0]} ${atom_j[1]} ${atom_j[2]})
            
            if (( $(echo "$dist < $min_distance" | bc -l) )); then
                min_distance=$dist
            fi
        done
    done
    
    # 最小距离阈值 (原始最小距离的0.8倍，约为1.6Å)
    local threshold=0.08  # 在Direct坐标中
    
    if (( $(echo "$min_distance < $threshold" | bc -l) )); then
        return 1
    else
        return 0
    fi
}

# 函数：生成随机构型
generate_random_config() {
    local config_id=$1
    local atoms=()
    
    # 读取原子位置
    while IFS= read -r line; do
        if [[ $line =~ ^[0-9] ]]; then
            atoms+=("$line")
        fi
    done <<< "$ORIGINAL_STRU"
    
    # 随机扰动参数
    local max_displacement=0.05  # 最大位移 (Direct坐标)
    local center_x=0.5
    local center_y=0.5
    
    # 生成随机构型
    local new_atoms=()
    for atom in "${atoms[@]}"; do
        local atom_data=($atom)
        local x=${atom_data[0]}
        local y=${atom_data[1]}
        local z=${atom_data[2]}
        
        # 随机位移
        local dx=$(echo "scale=10; ($RANDOM / 32767.0 - 0.5) * 2 * $max_displacement" | bc -l)
        local dy=$(echo "scale=10; ($RANDOM / 32767.0 - 0.5) * 2 * $max_displacement" | bc -l)
        
        # 应用位移
        local new_x=$(echo "scale=10; $x + $dx" | bc -l)
        local new_y=$(echo "scale=10; $y + $dy" | bc -l)
        
        # 处理周期性边界条件
        new_x=$(echo "scale=10; if ($new_x < 0) $new_x + 1 else if ($new_x >= 1) $new_x - 1 else $new_x" | bc -l)
        new_y=$(echo "scale=10; if ($new_y < 0) $new_y + 1 else if ($new_y >= 1) $new_y - 1 else $new_y" | bc -l)
        
        new_atoms+=("$new_x $new_y $z ${atom_data[3]} ${atom_data[4]} ${atom_data[5]} ${atom_data[6]} ${atom_data[7]}")
    done
    
    # 检查原子距离
    if check_atomic_distances "${new_atoms[@]}"; then
        printf '%s\n' "${new_atoms[@]}"
        return 0
    else
        return 1
    fi
}

# 函数：生成缩放构型
generate_scaled_config() {
    local config_id=$1
    local atoms=()
    
    # 读取原子位置
    while IFS= read -r line; do
        if [[ $line =~ ^[0-9] ]]; then
            atoms+=("$line")
        fi
    done <<< "$ORIGINAL_STRU"
    
    # 随机缩放因子
    local scale_factor=$(echo "scale=10; 0.85 + ($RANDOM / 32767.0) * 0.3" | bc -l)  # 0.85-1.15
    
    # 固定原子索引 (每层固定一个)
    local fixed_indices=(0 6 12 18)
    local center_x=0.5
    local center_y=0.5
    
    # 应用缩放
    local new_atoms=()
    for ((i=0; i<${#atoms[@]}; i++)); do
        local atom=(${atoms[i]})
        local x=${atom[0]} y=${atom[1]} z=${atom[2]}
        
        # 检查是否是固定原子
        local is_fixed=0
        for idx in "${fixed_indices[@]}"; do
            if [ $i -eq $idx ]; then
                is_fixed=1
                break
            fi
        done
        
        if [ $is_fixed -eq 1 ]; then
            # 固定原子，不缩放
            new_atoms+=("$x $y $z ${atom[3]} ${atom[4]} ${atom[5]} ${atom[6]} ${atom[7]}")
        else
            # 以超胞中心为参考点进行缩放
            local scaled_x=$(echo "scale=10; $center_x + ($x - $center_x) * $scale_factor" | bc -l)
            local scaled_y=$(echo "scale=10; $center_y + ($y - $center_y) * $scale_factor" | bc -l)
            
            # 处理周期性边界条件
            scaled_x=$(echo "scale=10; if ($scaled_x < 0) $scaled_x + 1 else if ($scaled_x >= 1) $scaled_x - 1 else $scaled_x" | bc -l)
            scaled_y=$(echo "scale=10; if ($scaled_y < 0) $scaled_y + 1 else if ($scaled_y >= 1) $scaled_y - 1 else $scaled_y" | bc -l)
            
            new_atoms+=("$scaled_x $scaled_y $z ${atom[3]} ${atom[4]} ${atom[5]} ${atom[6]} ${atom[7]}")
        fi
    done
    
    # 检查原子距离
    if check_atomic_distances "${new_atoms[@]}"; then
        printf '%s\n' "${new_atoms[@]}"
        return 0
    else
        return 1
    fi
}

# 函数：生成剪切构型
generate_shear_config() {
    local config_id=$1
    local atoms=()
    
    # 读取原子位置
    while IFS= read -r line; do
        if [[ $line =~ ^[0-9] ]]; then
            atoms+=("$line")
        fi
    done <<< "$ORIGINAL_STRU"
    
    # 随机剪切参数
    local shear_x=$(echo "scale=10; ($RANDOM / 32767.0 - 0.5) * 0.1" | bc -l)  # -0.05到0.05
    local shear_y=$(echo "scale=10; ($RANDOM / 32767.0 - 0.5) * 0.1" | bc -l)  # -0.05到0.05
    
    # 应用剪切变形
    local new_atoms=()
    for atom in "${atoms[@]}"; do
        local atom_data=($atom)
        local x=${atom_data[0]}
        local y=${atom_data[1]}
        local z=${atom_data[2]}
        
        # 剪切变形
        local new_x=$(echo "scale=10; $x + $shear_x * $y" | bc -l)
        local new_y=$(echo "scale=10; $y + $shear_y * $x" | bc -l)
        
        # 处理周期性边界条件
        new_x=$(echo "scale=10; if ($new_x < 0) $new_x + 1 else if ($new_x >= 1) $new_x - 1 else $new_x" | bc -l)
        new_y=$(echo "scale=10; if ($new_y < 0) $new_y + 1 else if ($new_y >= 1) $new_y - 1 else $new_y" | bc -l)
        
        new_atoms+=("$new_x $new_y $z ${atom_data[3]} ${atom_data[4]} ${atom_data[5]} ${atom_data[6]} ${atom_data[7]}")
    done
    
    # 检查原子距离
    if check_atomic_distances "${new_atoms[@]}"; then
        printf '%s\n' "${new_atoms[@]}"
        return 0
    else
        return 1
    fi
}

# 主循环
TOTAL_CONFIGS=1500
MAX_ATTEMPTS=1800  # 最大尝试次数，避免无限循环
completed_configs=0
failed_configs=0
attempt_count=0

echo "开始生成 $TOTAL_CONFIGS 个构型..."
echo "开始时间: $(date)"

while [ $completed_configs -lt $TOTAL_CONFIGS ] && [ $attempt_count -lt $MAX_ATTEMPTS ]; do
    attempt_count=$((attempt_count + 1))
    config_id=$((completed_configs + 1))
    config_name=$(printf "config_%04d" $config_id)
    config_dir="config_results_1500/$config_name"
    
    mkdir -p "$config_dir"
    cd "$config_dir"
    
    # 随机选择构型生成方法
    method=$((RANDOM % 3))
    case $method in
        0)
            new_atoms=$(generate_random_config $config_id)
            ;;
        1)
            new_atoms=$(generate_scaled_config $config_id)
            ;;
        2)
            new_atoms=$(generate_shear_config $config_id)
            ;;
    esac
    
    if [ $? -ne 0 ]; then
        echo "跳过 $config_name - 原子距离过近 (尝试 $attempt_count)"
        cd ../..
        failed_configs=$((failed_configs + 1))
        continue
    fi
    
    # 创建新的STRU文件
    cat > STRU << EOF
ATOMIC_SPECIES
Be    9.012  Be_ONCV_PBE-1.0.upf

NUMERICAL_ORBITAL
Be_gga_10au_100Ry_6s2p.orb

LATTICE_CONSTANT
1.889726

LATTICE_VECTORS
11.11980000000  0.000000000000  0.000000000000
0.000000000000  9.630000000000  0.000000000000
0.000000000000  0.000000000000  20.00000000000

ATOMIC_POSITIONS
Direct

Be
0.000
24
$new_atoms
EOF
    
    # 复制INPUT文件
    cp ../../INPUT_PARAMETERS INPUT
    
    # 运行ABACUS计算
    echo "计算构型: $config_name ($completed_configs/$TOTAL_CONFIGS)"
    OMP_NUM_THREADS=1 mpirun -n $CORE $ABACUS_CMD > out.log 2> err.log &
    pid=$!
    wait $pid
    
    # 提取结果
    if [ $? -eq 0 ] && [ -f "OUT.ABACUS/running_scf.log" ]; then
        ENERGY=$(grep "FINAL_ETOT_IS" OUT.ABACUS/running_scf.log | tail -1 | awk '{print $2}')
        VOLUME=$(grep "Volume" OUT.ABACUS/running_scf.log | awk '{print $4}')
        
        # 提取应力
        if grep -q "TOTAL-STRESS" OUT.ABACUS/running_scf.log; then
            STRESS_LINE=$(grep -A 4 "TOTAL-STRESS" OUT.ABACUS/running_scf.log | tail -3 | head -1)
            STRESS_XX=$(echo $STRESS_LINE | awk '{print $1}')
            STRESS_YY=$(echo $STRESS_LINE | awk '{print $2}')
            STRESS_ZZ=$(echo $STRESS_LINE | awk '{print $3}')
        else
            STRESS_XX="N/A"
            STRESS_YY="N/A"
            STRESS_ZZ="N/A"
        fi
        
        # 记录结果
        cd ../..
        echo "$config_name $ENERGY $VOLUME $STRESS_XX $STRESS_YY $STRESS_ZZ" >> $RESULT_FILE
        completed_configs=$((completed_configs + 1))
        
        # 显示进度
        progress=$(echo "scale=2; $completed_configs * 100 / $TOTAL_CONFIGS" | bc -l)
        echo "进度: $progress% - 完成: $config_name, 能量 = $ENERGY"
    else
        cd ../..
        echo "$config_name ERROR ERROR ERROR ERROR ERROR" >> $RESULT_FILE
        echo "错误: $config_name 计算失败"
        failed_configs=$((failed_configs + 1))
    fi
    
    # 每完成20个构型保存一次进度
    if [ $((completed_configs % 20)) -eq 0 ]; then
        echo "检查点: 已完成 $completed_configs/$TOTAL_CONFIGS 个构型"
        # 备份结果文件
        cp $RESULT_FILE "${RESULT_FILE}.backup"
    fi
done

echo "大规模构型计算完成!"
echo "完成时间: $(date)"
echo "成功计算: $completed_configs/$TOTAL_CONFIGS 个构型"
echo "失败/跳过: $failed_configs 个构型"
echo "总尝试次数: $attempt_count 次"
echo "成功率: $(echo "scale=2; $completed_configs * 100 / $attempt_count" | bc -l)%"
echo "结果保存在 $RESULT_FILE"
echo "所有构型数据保存在 config_results_1500/ 目录"

# 如果未达到目标数量，给出警告
if [ $completed_configs -lt $TOTAL_CONFIGS ]; then
    echo "警告: 只完成了 $completed_configs 个构型，未达到目标 $TOTAL_CONFIGS"
    echo "可能需要调整位移参数或增加最大尝试次数"
fi